import os
import shutil
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

def organize_dataset(data_root, output_root, test_size=0.1, val_size=0.1, random_seed=42):
    train_dir = os.path.join(output_root, 'train')
    val_dir = os.path.join(output_root, 'val')
    test_dir = os.path.join(output_root, 'test')
    
    for dir_path in [train_dir, val_dir, test_dir]:
        os.makedirs(os.path.join(dir_path, 'images'), exist_ok=True)
        os.makedirs(os.path.join(dir_path, 'labels'), exist_ok=True)
        os.makedirs(os.path.join(dir_path, 'metadata'), exist_ok=True)
    
    metadata_path = os.path.join(data_root, 'metadata', 'patch_metadata.csv')
    metadata_df = pd.read_csv(metadata_path)
    
    # 按public_id分组，确保同一个患者的所有骨折块在同一集合中!
    patient_ids = metadata_df['public_id'].unique()
    
    train_val_ids, test_ids = train_test_split(
        patient_ids, test_size=test_size, random_state=random_seed
    )
    
    train_ids, val_ids = train_test_split(
        train_val_ids, test_size=val_size/(1-test_size), random_state=random_seed
    )
    
    print(f"数据集划分:")
    print(f"  训练集患者: {len(train_ids)}")
    print(f"  验证集患者: {len(val_ids)}")
    print(f"  测试集患者: {len(test_ids)}")
    
    train_df = metadata_df[metadata_df['public_id'].isin(train_ids)]
    val_df = metadata_df[metadata_df['public_id'].isin(val_ids)]
    test_df = metadata_df[metadata_df['public_id'].isin(test_ids)]
    
    print(f"  训练集样本: {len(train_df)}")
    print(f"  验证集样本: {len(val_df)}")
    print(f"  测试集样本: {len(test_df)}")
    
    train_df.to_csv(os.path.join(train_dir, 'metadata', 'metadata.csv'), index=False)
    val_df.to_csv(os.path.join(val_dir, 'metadata', 'metadata.csv'), index=False)
    test_df.to_csv(os.path.join(test_dir, 'metadata', 'metadata.csv'), index=False)
    

    def copy_files(df, dest_dir):
        for _, row in df.iterrows():

            src_ct = os.path.join(data_root, row['ct_path'])
            src_mask = os.path.join(data_root, row['mask_path'])
            dest_ct = os.path.join(dest_dir, 'images', os.path.basename(row['ct_path']))
            dest_mask = os.path.join(dest_dir, 'labels', os.path.basename(row['mask_path']))
            shutil.copy2(src_ct, dest_ct)
            shutil.copy2(src_mask, dest_mask)
    
    print("复制训练集文件...")
    copy_files(train_df, train_dir)
    
    print("复制验证集文件...")
    copy_files(val_df, val_dir)
    
    print("复制测试集文件...")
    copy_files(test_df, test_dir)
    
    print("数据集组织完成!")
    return train_dir, val_dir, test_dir

if __name__ == "__main__":
    """划分数据集"""
    data_root = './ribfrac-patches/z_spacing'  # 原始数据目录
    output_root = './ribfrac-dataset/z_spacing'  # 输出目录
    

    train_dir, val_dir, test_dir = organize_dataset(
        data_root, output_root, test_size=0.1, val_size=0.1
    )
    
    print(f"\n数据集结构:")
    print(f"  训练集: {train_dir}")
    print(f"  验证集: {val_dir}")
    print(f"  测试集: {test_dir}")